{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载向量数据库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Created a chunk of size 109, which is longer than the specified 50\n",
      "Created a chunk of size 65, which is longer than the specified 50\n",
      "Created a chunk of size 143, which is longer than the specified 50\n",
      "Created a chunk of size 833, which is longer than the specified 50\n",
      "Created a chunk of size 263, which is longer than the specified 50\n",
      "Created a chunk of size 304, which is longer than the specified 50\n",
      "Created a chunk of size 609, which is longer than the specified 50\n",
      "Created a chunk of size 174, which is longer than the specified 50\n",
      "Created a chunk of size 449, which is longer than the specified 50\n",
      "Created a chunk of size 280, which is longer than the specified 50\n",
      "Created a chunk of size 325, which is longer than the specified 50\n",
      "Created a chunk of size 560, which is longer than the specified 50\n",
      "Created a chunk of size 318, which is longer than the specified 50\n",
      "Created a chunk of size 68, which is longer than the specified 50\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_text_splitters import CharacterTextSplitter\n",
    "\n",
    "from tongyi.embeddings import TongyiEmbeddings\n",
    "from langchain_community.vectorstores import FAISS\n",
    "\n",
    "# Load the document, split it into chunks, embed each chunk and load it into the vector store.\n",
    "raw_documents = TextLoader('./examples/rag.txt').load()\n",
    "text_splitter = CharacterTextSplitter(separator='\\n\\n\\n', chunk_size=50, chunk_overlap=4)\n",
    "documents = text_splitter.split_documents(raw_documents)\n",
    "db = FAISS.from_documents(documents, TongyiEmbeddings())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LCEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='一、在哪里可以了解高考成绩、志愿填报时间和方式、各高校招生计划、往年录取参考等志愿填报权威信息？\\n各省级教育行政部门或招生考试机构官方网站、微信公众号等权威渠道都会公布今年高考各阶段工作时间安排，包括高考成绩公布时间和查询方式、志愿填报时间，以及今年各高校招生计划、往年录取情况参考等权威信息。考生和家长要及时关注本地官方权威渠道发布的消息内容。\\n考生高考志愿是高校录取的重要依据，请广大考生务必按照省级招生考试机构相关要求按时完成志愿填报。前期，教育部已会同有关部门协调互联网平台对省级招生考试机构和高校的官方网站、微信公众号等进行了权威标识，请广大考生在信息查询时认准官方权威渠道，切勿轻信网络不实信息。', metadata={'source': './examples/rag.txt'}),\n",
       " Document(page_content='2024年高考是黑龙江、甘肃、吉林、安徽、江西、贵州、广西7个省份（中国第四批高考综合改革省份）的第一届落地实施的新高考。 [3]', metadata={'source': './examples/rag.txt'})]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retriever = db.as_retriever(search_type=\"similarity\", search_kwargs={\"k\": 2})\n",
    "\n",
    "retriever.invoke(\"哪里可以了解高考成绩\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MultiQueryRetriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='一、在哪里可以了解高考成绩、志愿填报时间和方式、各高校招生计划、往年录取参考等志愿填报权威信息？\\n各省级教育行政部门或招生考试机构官方网站、微信公众号等权威渠道都会公布今年高考各阶段工作时间安排，包括高考成绩公布时间和查询方式、志愿填报时间，以及今年各高校招生计划、往年录取情况参考等权威信息。考生和家长要及时关注本地官方权威渠道发布的消息内容。\\n考生高考志愿是高校录取的重要依据，请广大考生务必按照省级招生考试机构相关要求按时完成志愿填报。前期，教育部已会同有关部门协调互联网平台对省级招生考试机构和高校的官方网站、微信公众号等进行了权威标识，请广大考生在信息查询时认准官方权威渠道，切勿轻信网络不实信息。', metadata={'source': './examples/rag.txt'}),\n",
       " Document(page_content='三、高校招生章程有什么作用，如何查询？\\n高校招生章程由学校依据相关法律规定和国家招生政策制定，是学校开展招生工作的依据。考生在填报志愿前，应仔细查阅拟报考高校的招生章程，全面了解高校招生办法和相关招生要求。\\n主要查询途径有：中国高等教育学生信息网的“阳光高考”信息平台（https://gaokao.chsi.com.cn）；各高校官方招生网站等。', metadata={'source': './examples/rag.txt'}),\n",
       " Document(page_content='2024年高考是黑龙江、甘肃、吉林、安徽、江西、贵州、广西7个省份（中国第四批高考综合改革省份）的第一届落地实施的新高考。 [3]', metadata={'source': './examples/rag.txt'})]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
    "from langchain_community.chat_models import ChatTongyi\n",
    "\n",
    "llm = ChatTongyi()\n",
    "\n",
    "retriever_from_llm = MultiQueryRetriever.from_llm(\n",
    "    retriever=retriever, llm=llm, include_original=True\n",
    ")\n",
    "\n",
    "unique_docs = retriever_from_llm.invoke(\"哪里可以了解高考成绩\")\n",
    "unique_docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(unique_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Contextual compression\n",
    "## LLMChainExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='各省级教育行政部门或招生考试机构官方网站、微信公众号等权威渠道都会公布今年高考各阶段工作时间安排，包括高考成绩公布时间和查询方式、志愿填报时间，以及今年各高校招生计划、往年录取情况参考等权威信息。', metadata={'source': './examples/rag.txt'})]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import LLMChainExtractor\n",
    "\n",
    "compressor = LLMChainExtractor.from_llm(llm)\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=compressor, base_retriever=retriever\n",
    ")\n",
    "\n",
    "compressed_docs = compression_retriever.invoke(\n",
    "    \"哪里可以了解高考成绩\"\n",
    ")\n",
    "compressed_docs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLMChainFilter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='一、在哪里可以了解高考成绩、志愿填报时间和方式、各高校招生计划、往年录取参考等志愿填报权威信息？\\n各省级教育行政部门或招生考试机构官方网站、微信公众号等权威渠道都会公布今年高考各阶段工作时间安排，包括高考成绩公布时间和查询方式、志愿填报时间，以及今年各高校招生计划、往年录取情况参考等权威信息。考生和家长要及时关注本地官方权威渠道发布的消息内容。\\n考生高考志愿是高校录取的重要依据，请广大考生务必按照省级招生考试机构相关要求按时完成志愿填报。前期，教育部已会同有关部门协调互联网平台对省级招生考试机构和高校的官方网站、微信公众号等进行了权威标识，请广大考生在信息查询时认准官方权威渠道，切勿轻信网络不实信息。', metadata={'source': './examples/rag.txt'})]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.retrievers.document_compressors import LLMChainFilter\n",
    "\n",
    "_filter = LLMChainFilter.from_llm(llm)\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=_filter, base_retriever=retriever\n",
    ")\n",
    "\n",
    "compressed_docs = compression_retriever.invoke(\n",
    "    \"哪里可以了解高考成绩\"\n",
    ")\n",
    "compressed_docs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EmbeddingsFilter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.document_compressors import EmbeddingsFilter\n",
    "\n",
    "embeddings = TongyiEmbeddings()\n",
    "embeddings_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.76)\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=embeddings_filter, base_retriever=retriever\n",
    ")\n",
    "\n",
    "compressed_docs = compression_retriever.invoke(\n",
    "    \"What did the president say about Ketanji Jackson Brown\"\n",
    ")\n",
    "compressed_docs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DocumentCompressorPipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers.document_compressors import DocumentCompressorPipeline\n",
    "from langchain_community.document_transformers import EmbeddingsRedundantFilter\n",
    "\n",
    "redundant_filter = EmbeddingsRedundantFilter(embeddings=embeddings)\n",
    "relevant_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.76)\n",
    "pipeline_compressor = DocumentCompressorPipeline(\n",
    "    transformers=[redundant_filter, relevant_filter]\n",
    ")\n",
    "\n",
    "compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=pipeline_compressor, base_retriever=retriever\n",
    ")\n",
    "\n",
    "compressed_docs = compression_retriever.invoke(\n",
    "    \"What did the president say about Ketanji Jackson Brown\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ensemble Retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.retrievers import BM25Retriever\n",
    "\n",
    "doc_list = [doc.page_content for doc in documents]\n",
    "bm25_retriever = BM25Retriever.from_texts(\n",
    "    doc_list, metadatas=[{\"source\": f\"BM25\"}] * len(doc_list)\n",
    ")\n",
    "bm25_retriever.k = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='十、录取通知书何时能收到？\\n高校一般会在录取结束后一周左右向录取新生寄发录取通知书。若考生在省级招生考试机构或高校官方网站上查询到了录取结果，一直没有收到录取通知书，可及时联系录取高校公布的招生咨询电话查询本人录取通知书邮寄情况。', metadata={'source': 'BM25'}),\n",
       " Document(page_content='一、在哪里可以了解高考成绩、志愿填报时间和方式、各高校招生计划、往年录取参考等志愿填报权威信息？\\n各省级教育行政部门或招生考试机构官方网站、微信公众号等权威渠道都会公布今年高考各阶段工作时间安排，包括高考成绩公布时间和查询方式、志愿填报时间，以及今年各高校招生计划、往年录取情况参考等权威信息。考生和家长要及时关注本地官方权威渠道发布的消息内容。\\n考生高考志愿是高校录取的重要依据，请广大考生务必按照省级招生考试机构相关要求按时完成志愿填报。前期，教育部已会同有关部门协调互联网平台对省级招生考试机构和高校的官方网站、微信公众号等进行了权威标识，请广大考生在信息查询时认准官方权威渠道，切勿轻信网络不实信息。', metadata={'source': './examples/rag.txt'}),\n",
       " Document(page_content='九、录取工作采用什么方式，一般什么时间开始？\\n高校招生实行计算机远程网上录取，各省（区、市）录取工作一般于7月上旬开始，8月底之前结束。', metadata={'source': 'BM25'}),\n",
       " Document(page_content='2024年高考是黑龙江、甘肃、吉林、安徽、江西、贵州、广西7个省份（中国第四批高考综合改革省份）的第一届落地实施的新高考。 [3]', metadata={'source': './examples/rag.txt'})]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.retrievers import EnsembleRetriever\n",
    "\n",
    "# initialize the ensemble retriever\n",
    "ensemble_retriever = EnsembleRetriever(\n",
    "    retrievers=[bm25_retriever, retriever], weights=[0.5, 0.5]\n",
    ")\n",
    "\n",
    "docs = ensemble_retriever.invoke(\"哪里可以了解高考成绩\")\n",
    "docs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MultiVector Retriever\n",
    "## Smaller chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(15, 70)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.retrievers.multi_vector import MultiVectorRetriever\n",
    "from langchain.storage import InMemoryByteStore\n",
    "\n",
    "import uuid\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "doc_ids = [str(uuid.uuid4()) for _ in documents]\n",
    "id_key = \"doc_id\"\n",
    "\n",
    "# The splitter to use to create smaller chunks\n",
    "child_text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)\n",
    "sub_docs = []\n",
    "for i, doc in enumerate(documents):\n",
    "    _id = doc_ids[i]\n",
    "    _sub_docs = child_text_splitter.split_documents([doc])\n",
    "    for _doc in _sub_docs:\n",
    "        _doc.metadata[id_key] = _id\n",
    "    sub_docs.extend(_sub_docs)\n",
    "\n",
    "len(documents), len(sub_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The vectorstore to use to index the child chunks\n",
    "vectorstore = FAISS.from_documents(sub_docs, TongyiEmbeddings())\n",
    "# The storage layer for the parent documents\n",
    "store = InMemoryByteStore()\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    byte_store=store,\n",
    "    id_key=id_key,\n",
    ")\n",
    "retriever.docstore.mset(list(zip(doc_ids, documents)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(page_content='一、在哪里可以了解高考成绩、志愿填报时间和方式、各高校招生计划、往年录取参考等志愿填报权威信息？', metadata={'source': './examples/rag.txt', 'doc_id': 'd855063c-ad52-4a09-a304-aa9d2b2ebd17'})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Vectorstore alone retrieves the small chunks\n",
    "retriever.vectorstore.similarity_search(\"哪里可以了解高考成绩\")[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(page_content='一、在哪里可以了解高考成绩、志愿填报时间和方式、各高校招生计划、往年录取参考等志愿填报权威信息？\\n各省级教育行政部门或招生考试机构官方网站、微信公众号等权威渠道都会公布今年高考各阶段工作时间安排，包括高考成绩公布时间和查询方式、志愿填报时间，以及今年各高校招生计划、往年录取情况参考等权威信息。考生和家长要及时关注本地官方权威渠道发布的消息内容。\\n考生高考志愿是高校录取的重要依据，请广大考生务必按照省级招生考试机构相关要求按时完成志愿填报。前期，教育部已会同有关部门协调互联网平台对省级招生考试机构和高校的官方网站、微信公众号等进行了权威标识，请广大考生在信息查询时认准官方权威渠道，切勿轻信网络不实信息。', metadata={'source': './examples/rag.txt'})"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Retriever returns larger chunks\n",
    "retriever.invoke(\"哪里可以了解高考成绩\")[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='2024年普通高等学校招生全国统一考试（简称：2024年全国高考），是中华人民共和国合格的高中毕业生或具有同等学力的考生参加的选拔性考试', metadata={'source': './examples/rag.txt', 'doc_id': 'ffccfda6-67bb-4454-b4a1-c81c011f1512'}),\n",
       " Document(page_content='[1-2]。2024年报名人数1342万人，比2023年增加51万人 [21]。', metadata={'source': './examples/rag.txt', 'doc_id': 'ffccfda6-67bb-4454-b4a1-c81c011f1512'}),\n",
       " Document(page_content='2024年高考是黑龙江、甘肃、吉林、安徽、江西、贵州、广西7个省份（中国第四批高考综合改革省份）的第一届落地实施的新高考。 [3]', metadata={'source': './examples/rag.txt', 'doc_id': 'f64eaa30-9981-4433-96ed-2c454a9249ce'}),\n",
       " Document(page_content='2024年高考全国统考于2024年6月7日开始举行，部分省份考试时间为2天，实行新高考的省份为3-4天 [29]。5月31日，2024年高考试卷从北京发往全国', metadata={'source': './examples/rag.txt', 'doc_id': '43267b94-dff0-4524-a210-ad70dcee0ed8'}),\n",
       " Document(page_content='[22]。6月5日，2024年高考举报电话已开通，教育部教育考试院的举报电话为：010-62790357 [53] [77]。', metadata={'source': './examples/rag.txt', 'doc_id': '43267b94-dff0-4524-a210-ad70dcee0ed8'})]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub_docs[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parent Document Retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.retrievers import ParentDocumentRetriever\n",
    "from langchain_chroma import Chroma\n",
    "from langchain.storage import InMemoryStore\n",
    "\n",
    "# This text splitter is used to create the child documents\n",
    "child_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)\n",
    "# The vectorstore to use to index the child chunks (empty to start)\n",
    "# FAISS not supports empty initialzation\n",
    "vectorstore = Chroma(\n",
    "    collection_name=\"full_documents\", embedding_function=TongyiEmbeddings()\n",
    ")\n",
    "# The storage layer for the parent documents\n",
    "store = InMemoryStore()\n",
    "retriever = ParentDocumentRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    docstore=store,\n",
    "    child_splitter=child_splitter,\n",
    ")\n",
    "\n",
    "retriever.add_documents(documents, ids=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sub_docs = vectorstore.similarity_search(\"哪里可以了解高考成绩\")\n",
    "sub_docs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieved_docs = retriever.invoke(\"哪里可以了解高考成绩\")\n",
    "retrieved_docs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['2024年全国高考是中国的一项重要考试，用于选拔高中毕业生和具备同等学历的考生，2024年的报名人数达到了1342万人，相比上一年增长了51万人。',\n",
       " '2024年，中国有7个省份（黑龙江、甘肃、吉林、安徽、江西、贵州和广西）将首次实施新的高考制度，作为第四批改革省份。']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import uuid\n",
    "\n",
    "from langchain_core.documents import Document\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_community.chat_models import ChatTongyi\n",
    "\n",
    "chain = (\n",
    "    {\"doc\": lambda x: x.page_content}\n",
    "    | ChatPromptTemplate.from_template(\"概括以下内容:\\n\\n{doc}\")\n",
    "    | ChatTongyi(max_retries=0)\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "summaries = chain.batch(documents, {\"max_concurrency\": 5})\n",
    "summaries[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_ids = [str(uuid.uuid4()) for _ in documents]\n",
    "id_key = \"doc_id\"\n",
    "\n",
    "summary_docs = [\n",
    "    Document(page_content=s, metadata={id_key: doc_ids[i]})\n",
    "    for i, s in enumerate(summaries)\n",
    "]\n",
    "\n",
    "# The vectorstore to use to index the child chunks\n",
    "vectorstore = FAISS.from_documents(summary_docs, TongyiEmbeddings())\n",
    "# The storage layer for the parent documents\n",
    "store = InMemoryByteStore()\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    byte_store=store,\n",
    "    id_key=id_key,\n",
    ")\n",
    "retriever.docstore.mset(list(zip(doc_ids, documents)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(page_content='高考相关信息可以在各省级教育行政部门或招生考试机构的官方网站、微信公众号等权威渠道获取，包括成绩公布时间、查询方式、志愿填报时间、高校招生计划和历年录取参考等。考生和家长需密切关注官方发布的信息。志愿填报至关重要，考生必须遵循省级招生考试机构的规定。教育部已与相关部门合作确保官方渠道的权威性，提醒大家在查询时要识别官方标识，避免相信非官方的不实信息。', metadata={'doc_id': 'd5f6fbc3-3425-4c3e-914a-f10669c9ae53'})"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub_docs = vectorstore.similarity_search(\"哪里可以了解高考成绩\")\n",
    "sub_docs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(page_content='一、在哪里可以了解高考成绩、志愿填报时间和方式、各高校招生计划、往年录取参考等志愿填报权威信息？\\n各省级教育行政部门或招生考试机构官方网站、微信公众号等权威渠道都会公布今年高考各阶段工作时间安排，包括高考成绩公布时间和查询方式、志愿填报时间，以及今年各高校招生计划、往年录取情况参考等权威信息。考生和家长要及时关注本地官方权威渠道发布的消息内容。\\n考生高考志愿是高校录取的重要依据，请广大考生务必按照省级招生考试机构相关要求按时完成志愿填报。前期，教育部已会同有关部门协调互联网平台对省级招生考试机构和高校的官方网站、微信公众号等进行了权威标识，请广大考生在信息查询时认准官方权威渠道，切勿轻信网络不实信息。', metadata={'source': './examples/rag.txt'})"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retrieved_docs = retriever.invoke(\"哪里可以了解高考成绩\")\n",
    "retrieved_docs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hypothetical Queries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1. 2024年全国高考的全称是什么？',\n",
       " '2. 与2023年相比，2024年全国高考的报名人数有何变化？',\n",
       " '3. 能否提供2023年全国高考的报名人数数据作为对比？']"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import AIMessage\n",
    "from langchain_core.exceptions import OutputParserException\n",
    "\n",
    "def custom_parse(ai_message: AIMessage) -> str:\n",
    "    \"\"\"Parse the AI message.\"\"\"\n",
    "    if '\\n\\n' in ai_message.content:\n",
    "        return ai_message.content.split('\\n\\n')\n",
    "    elif '\\n' in ai_message.content:\n",
    "        return ai_message.content.split('\\n')\n",
    "    else:\n",
    "        raise OutputParserException(\"Badly formed question!\")\n",
    "\n",
    "chain = (\n",
    "    {\"doc\": lambda x: x.page_content}\n",
    "    | ChatPromptTemplate.from_template(\"为下面内容生成3个合适的提问问题：\\n\\n{doc}\\n\\n#限制\\n生成的3个问题使用两个换行符，即```\\n\\n```符号进行隔开\")\n",
    "    | ChatTongyi(max_retries=0)\n",
    "    | custom_parse\n",
    ")\n",
    "\n",
    "hypothetical_questions = chain.batch(documents, {\"max_concurrency\": 5})\n",
    "hypothetical_questions[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_ids = [str(uuid.uuid4()) for _ in documents]\n",
    "id_key = \"doc_id\"\n",
    "\n",
    "question_docs = []\n",
    "for i, question_list in enumerate(hypothetical_questions):\n",
    "    question_docs.extend(\n",
    "        [Document(page_content=s, metadata={id_key: doc_ids[i]}) for s in question_list]\n",
    "    )\n",
    "\n",
    "# The vectorstore to use to index the child chunks\n",
    "vectorstore = FAISS.from_documents(question_docs, TongyiEmbeddings())\n",
    "# The storage layer for the parent documents\n",
    "store = InMemoryByteStore()\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    byte_store=store,\n",
    "    id_key=id_key,\n",
    ")\n",
    "retriever.docstore.mset(list(zip(doc_ids, documents)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(page_content='1. 从哪里可以获得高考成绩查询的具体时间和方式，以及志愿填报的详细指导？', metadata={'doc_id': '120780d1-f9f2-4ee4-ac03-d806343878b8'})"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sub_docs = vectorstore.similarity_search(\"哪里可以了解高考成绩\")\n",
    "sub_docs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(page_content='一、在哪里可以了解高考成绩、志愿填报时间和方式、各高校招生计划、往年录取参考等志愿填报权威信息？\\n各省级教育行政部门或招生考试机构官方网站、微信公众号等权威渠道都会公布今年高考各阶段工作时间安排，包括高考成绩公布时间和查询方式、志愿填报时间，以及今年各高校招生计划、往年录取情况参考等权威信息。考生和家长要及时关注本地官方权威渠道发布的消息内容。\\n考生高考志愿是高校录取的重要依据，请广大考生务必按照省级招生考试机构相关要求按时完成志愿填报。前期，教育部已会同有关部门协调互联网平台对省级招生考试机构和高校的官方网站、微信公众号等进行了权威标识，请广大考生在信息查询时认准官方权威渠道，切勿轻信网络不实信息。', metadata={'source': './examples/rag.txt'})"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "retrieved_docs = retriever.invoke(\"哪里可以了解高考成绩\")\n",
    "retrieved_docs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Self-querying"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_chroma import Chroma\n",
    "from langchain_core.documents import Document\n",
    "\n",
    "docs = [\n",
    "    Document(\n",
    "        page_content=\"A bunch of scientists bring back dinosaurs and mayhem breaks loose\",\n",
    "        metadata={\"year\": 1993, \"rating\": 7.7, \"genre\": \"science fiction\"},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Leo DiCaprio gets lost in a dream within a dream within a dream within a ...\",\n",
    "        metadata={\"year\": 2010, \"director\": \"Christopher Nolan\", \"rating\": 8.2},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea\",\n",
    "        metadata={\"year\": 2006, \"director\": \"Satoshi Kon\", \"rating\": 8.6},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"A bunch of normal-sized women are supremely wholesome and some men pine after them\",\n",
    "        metadata={\"year\": 2019, \"director\": \"Greta Gerwig\", \"rating\": 8.3},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Toys come alive and have a blast doing so\",\n",
    "        metadata={\"year\": 1995, \"genre\": \"animated\"},\n",
    "    ),\n",
    "    Document(\n",
    "        page_content=\"Three men walk into the Zone, three men walk out of the Zone\",\n",
    "        metadata={\n",
    "            \"year\": 1979,\n",
    "            \"director\": \"Andrei Tarkovsky\",\n",
    "            \"genre\": \"thriller\",\n",
    "            \"rating\": 9.9,\n",
    "        },\n",
    "    ),\n",
    "]\n",
    "vectorstore = Chroma.from_documents(docs, TongyiEmbeddings())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.query_constructor.base import AttributeInfo\n",
    "from langchain.retrievers.self_query.base import SelfQueryRetriever\n",
    "\n",
    "metadata_field_info = [\n",
    "    AttributeInfo(\n",
    "        name=\"genre\",\n",
    "        description=\"The genre of the movie. One of ['science fiction', 'comedy', 'drama', 'thriller', 'romance', 'action', 'animated']\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"year\",\n",
    "        description=\"The year the movie was released\",\n",
    "        type=\"integer\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"director\",\n",
    "        description=\"The name of the movie director\",\n",
    "        type=\"string\",\n",
    "    ),\n",
    "    AttributeInfo(\n",
    "        name=\"rating\", description=\"A 1-10 rating for the movie\", type=\"float\"\n",
    "    ),\n",
    "]\n",
    "document_content_description = \"Brief summary of a movie\"\n",
    "llm = ChatTongyi()\n",
    "retriever = SelfQueryRetriever.from_llm(\n",
    "    llm,\n",
    "    vectorstore,\n",
    "    document_content_description,\n",
    "    metadata_field_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This example only specifies a filter\n",
    "retriever.invoke(\"I want to watch a movie rated higher than 8.5\")\n",
    "\"\"\"\n",
    "[Document(page_content='Three men walk into the Zone, three men walk out of the Zone', metadata={'director': 'Andrei Tarkovsky', 'genre': 'thriller', 'rating': 9.9, 'year': 1979}),\n",
    " Document(page_content='A psychologist / detective gets lost in a series of dreams within dreams within dreams and Inception reused the idea', metadata={'director': 'Satoshi Kon', 'rating': 8.6, 'year': 2006})]\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 原理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chains.query_constructor.base import (\n",
    "    StructuredQueryOutputParser,\n",
    "    get_query_constructor_prompt,\n",
    ")\n",
    "from langchain.retrievers.self_query.chroma import ChromaTranslator\n",
    "from langchain_community.chat_models import ChatTongyi\n",
    "\n",
    "llm = ChatTongyi()\n",
    "\n",
    "prompt = get_query_constructor_prompt(\n",
    "    document_content_description,\n",
    "    metadata_field_info,\n",
    "    allowed_comparators=ChromaTranslator.allowed_comparators,\n",
    "    allowed_operators=ChromaTranslator.allowed_operators\n",
    ")\n",
    "output_parser = StructuredQueryOutputParser.from_components()\n",
    "query_constructor = prompt | llm | output_parser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<Comparator.EQ: 'eq'>,\n",
       " <Comparator.NE: 'ne'>,\n",
       " <Comparator.GT: 'gt'>,\n",
       " <Comparator.GTE: 'gte'>,\n",
       " <Comparator.LT: 'lt'>,\n",
       " <Comparator.LTE: 'lte'>]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ChromaTranslator.allowed_comparators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<Operator.AND: 'and'>, <Operator.OR: 'or'>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ChromaTranslator.allowed_operators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Your goal is to structure the user\\'s query to match the request schema provided below.\\n\\n<< Structured Request Schema >>\\nWhen responding use a markdown code snippet with a JSON object formatted in the following schema:\\n\\n```json\\n{\\n    \"query\": string \\\\ text string to compare to document contents\\n    \"filter\": string \\\\ logical condition statement for filtering documents\\n}\\n```\\n\\nThe query string should contain only text that is expected to match the contents of documents. Any conditions in the filter should not be mentioned in the query as well.\\n\\nA logical condition statement is composed of one or more comparison and logical operation statements.\\n\\nA comparison statement takes the form: `comp(attr, val)`:\\n- `comp` (eq | ne | gt | gte | lt | lte): comparator\\n- `attr` (string):  name of attribute to apply the comparison to\\n- `val` (string): is the comparison value\\n\\nA logical operation statement takes the form `op(statement1, statement2, ...)`:\\n- `op` (and | or): logical operator\\n- `statement1`, `statement2`, ... (comparison statements or logical operation statements): one or more statements to apply the operation to\\n\\nMake sure that you only use the comparators and logical operators listed above and no others.\\nMake sure that filters only refer to attributes that exist in the data source.\\nMake sure that filters only use the attributed names with its function names if there are functions applied on them.\\nMake sure that filters only use format `YYYY-MM-DD` when handling date data typed values.\\nMake sure that filters take into account the descriptions of attributes and only make comparisons that are feasible given the type of data being stored.\\nMake sure that filters are only used as needed. If there are no filters that should be applied return \"NO_FILTER\" for the filter value.\\n\\n<< Example 1. >>\\nData Source:\\n```json\\n{\\n    \"content\": \"Lyrics of a song\",\\n    \"attributes\": {\\n        \"artist\": {\\n            \"type\": \"string\",\\n            \"description\": \"Name of the song artist\"\\n        },\\n        \"length\": {\\n            \"type\": \"integer\",\\n            \"description\": \"Length of the song in seconds\"\\n        },\\n        \"genre\": {\\n            \"type\": \"string\",\\n            \"description\": \"The song genre, one of \"pop\", \"rock\" or \"rap\"\"\\n        }\\n    }\\n}\\n```\\n\\nUser Query:\\nWhat are songs by Taylor Swift or Katy Perry about teenage romance under 3 minutes long in the dance pop genre\\n\\nStructured Request:\\n```json\\n{\\n    \"query\": \"teenager love\",\\n    \"filter\": \"and(or(eq(\\\\\"artist\\\\\", \\\\\"Taylor Swift\\\\\"), eq(\\\\\"artist\\\\\", \\\\\"Katy Perry\\\\\")), lt(\\\\\"length\\\\\", 180), eq(\\\\\"genre\\\\\", \\\\\"pop\\\\\"))\"\\n}\\n```\\n\\n\\n<< Example 2. >>\\nData Source:\\n```json\\n{\\n    \"content\": \"Lyrics of a song\",\\n    \"attributes\": {\\n        \"artist\": {\\n            \"type\": \"string\",\\n            \"description\": \"Name of the song artist\"\\n        },\\n        \"length\": {\\n            \"type\": \"integer\",\\n            \"description\": \"Length of the song in seconds\"\\n        },\\n        \"genre\": {\\n            \"type\": \"string\",\\n            \"description\": \"The song genre, one of \"pop\", \"rock\" or \"rap\"\"\\n        }\\n    }\\n}\\n```\\n\\nUser Query:\\nWhat are songs that were not published on Spotify\\n\\nStructured Request:\\n```json\\n{\\n    \"query\": \"\",\\n    \"filter\": \"NO_FILTER\"\\n}\\n```\\n\\n\\n<< Example 3. >>\\nData Source:\\n```json\\n{\\n    \"content\": \"Brief summary of a movie\",\\n    \"attributes\": {\\n    \"genre\": {\\n        \"description\": \"The genre of the movie. One of [\\'science fiction\\', \\'comedy\\', \\'drama\\', \\'thriller\\', \\'romance\\', \\'action\\', \\'animated\\']\",\\n        \"type\": \"string\"\\n    },\\n    \"year\": {\\n        \"description\": \"The year the movie was released\",\\n        \"type\": \"integer\"\\n    },\\n    \"director\": {\\n        \"description\": \"The name of the movie director\",\\n        \"type\": \"string\"\\n    },\\n    \"rating\": {\\n        \"description\": \"A 1-10 rating for the movie\",\\n        \"type\": \"float\"\\n    }\\n}\\n}\\n```\\n\\nUser Query:\\ndummy question\\n\\nStructured Request:\\n'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt.format(query=\"dummy question\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StructuredQuery(query='dinosaurs', filter=Comparison(comparator=<Comparator.GT: 'gt'>, attribute='rating', value=8.5), limit=None)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_constructor.invoke(\"I want to watch a movie about dinosaurs rated higher than 8.5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2024高考RAG应用\n",
    "## without rag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='对不起，我无法提供具体的2024年高考报名人数信息，因为这些数据通常由各省份的教育考试机构或政府部门发布，而且会在高考报名开始前公布。对于这类实时数据，建议你关注当地教育部门或考试院的官方通知，或者在高考报名开始时查询相关公告。', response_metadata={'model_name': 'qwen-turbo', 'finish_reason': 'stop', 'request_id': '2cffd503-04a3-96e6-bed6-efd912311105', 'token_usage': {'input_tokens': 16, 'output_tokens': 67, 'total_tokens': 83}}, id='run-293addeb-4aa2-4ea9-b53a-861483c0114c-0')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_community.chat_models import ChatTongyi\n",
    "\n",
    "chat = ChatTongyi()\n",
    "\n",
    "chat.invoke('2024年高考报名人数是多少')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='高考报名时间每年可能会有所变动，具体以官方发布的通知为准。一般来说，广东省的高考报名时间通常在每年的11月份进行，持续一周左右。建议你关注广东省教育考试院或当地教育局的官方网站，他们会发布最准确的高考报名通知和时间安排。同时，也要注意报名截止日期，不要错过报名时间。', response_metadata={'model_name': 'qwen-turbo', 'finish_reason': 'stop', 'request_id': 'a2e34ac6-4188-9626-a93d-cfea2347401d', 'token_usage': {'input_tokens': 20, 'output_tokens': 74, 'total_tokens': 94}}, id='run-5a200266-3e07-4271-9f7e-a74f4b7156d9-0')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat.invoke('2024年高考，广东的报名时间是什么时候')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## with rag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Created a chunk of size 109, which is longer than the specified 50\n",
      "Created a chunk of size 65, which is longer than the specified 50\n",
      "Created a chunk of size 143, which is longer than the specified 50\n",
      "Created a chunk of size 833, which is longer than the specified 50\n",
      "Created a chunk of size 263, which is longer than the specified 50\n",
      "Created a chunk of size 304, which is longer than the specified 50\n",
      "Created a chunk of size 609, which is longer than the specified 50\n",
      "Created a chunk of size 174, which is longer than the specified 50\n",
      "Created a chunk of size 449, which is longer than the specified 50\n",
      "Created a chunk of size 280, which is longer than the specified 50\n",
      "Created a chunk of size 325, which is longer than the specified 50\n",
      "Created a chunk of size 560, which is longer than the specified 50\n",
      "Created a chunk of size 318, which is longer than the specified 50\n",
      "Created a chunk of size 68, which is longer than the specified 50\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_text_splitters import CharacterTextSplitter\n",
    "\n",
    "from tongyi.embeddings import TongyiEmbeddings\n",
    "from langchain_community.vectorstores import FAISS\n",
    "\n",
    "# Load the document, split it into chunks, embed each chunk and load it into the vector store.\n",
    "raw_documents = TextLoader('./examples/rag.txt').load()\n",
    "text_splitter = CharacterTextSplitter(separator='\\n\\n\\n', chunk_size=50, chunk_overlap=4)\n",
    "documents = text_splitter.split_documents(raw_documents)\n",
    "db = FAISS.from_documents(documents, TongyiEmbeddings())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1. 2024年全国高考的全称是什么？',\n",
       " '2. 与2023年相比，2024年全国高考的报名人数有何变化？',\n",
       " '3. 能否提供2023年全国高考的报名人数数据作为对比？']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.messages import AIMessage\n",
    "from langchain_core.exceptions import OutputParserException\n",
    "\n",
    "def custom_parse(ai_message: AIMessage) -> str:\n",
    "    \"\"\"Parse the AI message.\"\"\"\n",
    "    if '\\n\\n' in ai_message.content:\n",
    "        return ai_message.content.split('\\n\\n')\n",
    "    elif '\\n' in ai_message.content:\n",
    "        return ai_message.content.split('\\n')\n",
    "    else:\n",
    "        raise OutputParserException(\"Badly formed question!\")\n",
    "\n",
    "chain = (\n",
    "    {\"doc\": lambda x: x.page_content}\n",
    "    | ChatPromptTemplate.from_template(\"为下面内容生成3个合适的提问问题：\\n\\n{doc}\\n\\n#限制\\n生成的3个问题使用两个换行符，即```\\n\\n```符号进行隔开\")\n",
    "    | ChatTongyi(max_retries=0)\n",
    "    | custom_parse\n",
    ")\n",
    "\n",
    "hypothetical_questions = chain.batch(documents, {\"max_concurrency\": 5})\n",
    "hypothetical_questions[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "from langchain_core.documents import Document\n",
    "from langchain.storage import InMemoryByteStore\n",
    "from langchain.retrievers.multi_vector import MultiVectorRetriever\n",
    "\n",
    "doc_ids = [str(uuid.uuid4()) for _ in documents]\n",
    "id_key = \"doc_id\"\n",
    "\n",
    "question_docs = []\n",
    "for i, question_list in enumerate(hypothetical_questions):\n",
    "    question_docs.extend(\n",
    "        [Document(page_content=s, metadata={id_key: doc_ids[i]}) for s in question_list]\n",
    "    )\n",
    "\n",
    "# The vectorstore to use to index the child chunks\n",
    "vectorstore = FAISS.from_documents(question_docs, TongyiEmbeddings())\n",
    "# The storage layer for the parent documents\n",
    "store = InMemoryByteStore()\n",
    "retriever = MultiVectorRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    byte_store=store,\n",
    "    id_key=id_key,\n",
    ")\n",
    "retriever.docstore.mset(list(zip(doc_ids, documents)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2024年全国高考的报名人数达到1342万人。'"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "template = \"\"\"回答用户的问题，下面的内容可以作为你的知识依据：\n",
    "```\n",
    "{context}\n",
    "```\n",
    "\n",
    "用户的问题：{query}\n",
    "\"\"\"\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join([d.page_content for d in docs])\n",
    "\n",
    "\n",
    "rag_chain = (\n",
    "    {\"context\": retriever | format_docs, \"query\": RunnablePassthrough()}\n",
    "    | prompt\n",
    "    | chat\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "rag_chain.invoke(\"2024年高考报名人数是多少\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2024年广东的高考报名时间是2023年11月1日至10日。'"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rag_chain.invoke(\"2024年高考，广东的报名时间是什么时候\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
